//
// Created by lqk on 2022/1/30.
//

#include "ConvexMPController.h"

ConvexMPController::ConvexMPController(float dt, int loopsPerMPC) :
        _loopsPerMPC(loopsPerMPC),
        _nSegment(14),
        _dt(dt),
        _trotting(_nSegment, Eigen::Vector4i(0, 7, 7, 0), 7, "Trotting"),
        _standing(_nSegment,Eigen::Vector4i(0,0,0,0),14,"Standing")
{
    _dtMPC= dt * loopsPerMPC;
    for (int i = 0; i < 4; i++){
        _firstSwing[i] = true;
        _footForce[i].setZero();
        _posFootWorld[i].setZero();
    }


}
void ConvexMPController::Update(EstimatedState *const eState, LegCommand *const legsCommand,LegState *const legsState,QuadrupedKinematics*const quadKine) {
    bool omniMode = false;
    _rollDes=0.;
    _pitchDes=0.;
    _yawDes=eState->m_rpy[2]+_dt*_yawVelDes;
    OffsetDurationGait* gait=&_trotting;
    _currentGait=GaitName::trotting;
    gait->UpdatePhase(_loopsPerMPC, _loopsCounter);
    float swingDuration=gait->GetSwingDuration(_dtMPC);
    float stanceDuration=gait->GetStanceDuration(_dtMPC);
    Eigen::Vector3f velDesBody(_xVelDes,_yVelDes,0);
    Eigen::Vector3f velDesWorld=omniMode ? velDesBody : eState->m_rotation*velDesBody;
    Eigen::Vector3f velRobot=eState->m_velWorld;

    for(int i=0;i<4;i++){
        _posFootWorld[i]= eState->m_position + eState->m_rotation * (quadKine->GetHipLocation(i) + legsState[i].m_footPos);
        std::cout<<"Foot "<<i<<" pos:"<<_posFootWorld[i].transpose()<<"\n";
    }

    if(gait!=&_standing){
        _posDes+=_dt*Eigen::Vector3f(velDesWorld[0],velDesWorld[1],0);
    }
    if(_firstRun){
        for(int time=0;time<1000;time++){
            for(int leg=0;leg<4;leg++){
                Eigen::Vector3f q_des(0,-1,2);
                Eigen::Vector3f q_next=(q_des-legsState[leg].m_q)*0.001*time;
                legsCommand[leg].m_qDesired=q_next;
                legsCommand[leg].m_kdJoints(0,0)=1000;
                legsCommand[leg].m_kdJoints(1,1)=1000;
                legsCommand[leg].m_kdJoints(2,2)=1000;
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            }
        }
        for(int leg=0;leg<4;leg++){
            legsCommand[leg].m_kdJoints(0,0)=0;
            legsCommand[leg].m_kdJoints(1,1)=0;
            legsCommand[leg].m_kdJoints(2,2)=0;
        }

        _posDes[0]=eState->m_position[0];
        _posDes[1]=eState->m_position[1];
        _posDes[2]=eState->m_rpy[2];
        for(int i=0;i<4;i++){
            _swingPlanner[i].SetHeight(0.06);
            _swingPlanner[i].SetInitialPosition(_posFootWorld[i]);
            _swingPlanner[i].SetFinalPosition(_posFootWorld[i]);
//            std::cout<<"initial foot pos"<<_posFootWorld[i].transpose()<<std::endl;
        }
//        _swingPlanner[0].SetInitialPosition(Eigen::Vector3f (0.178219,-0.111,0.));
//        _swingPlanner[1].SetInitialPosition(Eigen::Vector3f (0.178219,0.111,0.));
//        _swingPlanner[2].SetInitialPosition(Eigen::Vector3f (-0.201781,-0.111,0.));
//        _swingPlanner[3].SetInitialPosition(Eigen::Vector3f (-0.201781,0.111,0.));
//        _swingPlanner[0].SetFinalPosition(Eigen::Vector3f (0.178219,-0.111,0.));
//        _swingPlanner[1].SetFinalPosition(Eigen::Vector3f (0.178219,0.111,0.));
//        _swingPlanner[2].SetFinalPosition(Eigen::Vector3f (-0.201781,-0.111,0.));
//        _swingPlanner[3].SetFinalPosition(Eigen::Vector3f (-0.201781,0.111,0.));
        _firstRun=false;

    }

    float side_sign[4] = {-1, 1, -1, 1};
    // * Calculate footstep location and set the swing trajectory
    for(int i=0;i<4;i++){
        if(_firstSwing[i]){
            _swingTimeRemain[i]=swingDuration;
        }else{
            _swingTimeRemain[i]-=_dt;
        }
        _swingPlanner[i].SetHeight(0.06);
        Eigen::Vector3f offset(0,side_sign[i]*.062,0);
        Eigen::Vector3f hippBody=(quadKine->GetHipLocation(i)+offset);
        std::cout<<"Hipp "<<i<<" location"<<hippBody.transpose()<<std::endl;
        Eigen::Vector3f hippWorldRot = Eigen::AngleAxisf(-_yawVelDes * stanceDuration/2., Eigen::Vector3f::UnitZ())*hippBody;
        Eigen::Vector3f pfHipPrimer = eState->m_position + eState->m_rotation * (hippWorldRot + velDesBody * _swingTimeRemain[i]);
        float pfxHeuristic = eState->m_velWorld[0] * stanceDuration * 0.5+0.05*(eState->m_velWorld[0]-velDesBody[0]);
        float pfyHeuristic = eState->m_velWorld[1] * stanceDuration * 0.5+0.05*(eState->m_velWorld[1]-velDesBody[1]);
        Eigen::Vector3f pfHeuristic(fminf(fmaxf(pfxHeuristic,-0.3),0.3),
                                    fminf(fmaxf(pfyHeuristic,-0.3),0.3),
                                    0);
        Eigen::Vector3f pf = pfHipPrimer + pfHeuristic;
        pf[2] = -0.002;
        _swingPlanner[i].SetFinalPosition(pf);
    }

    _loopsCounter++;
    // load LCM leg swing gains
    _kp << 700, 0, 0, 0, 700, 0, 0, 0, 2000;
    _kpStance = 0.0*_kp;
    _kd << 10, 0, 0, 0, 10, 0, 0, 0, 10;
    _kdStance = 1.0 * _kd;
    Eigen::Vector4f stancePhase = gait->GetStancePhase();
    Eigen::Vector4f swingPhase = gait->GetSwingPhase();
    int* mpcTable = gait->GetMPCTable();
    UpdateMPCIfNeeded(mpcTable,eState,omniMode);
    // * Update LegsCommand
    for(int foot=0;foot<4;foot++){
        // * Foot in swing state
        if(swingPhase[foot]>0){
//            std::cout<<"swing"<<std::endl;
            if(_firstSwing[foot]){
                _firstSwing[foot]= false;
                _swingPlanner[foot].SetInitialPosition(_posFootWorld[foot]);
            }
            _swingPlanner[foot].SetSwingPeriod(swingDuration);
            _swingPlanner[foot].Compute(swingPhase[foot]);
            Eigen::Vector3f posDesFootWorld = _swingPlanner[foot].GetPosition();
//            std::cout<<posDesFootWorld.transpose()<<std::endl;
            Eigen::Vector3f velDesFootWorld = _swingPlanner[foot].GetVelocity();
            Eigen::Vector3f posDesFootLeg = eState->m_rotation.transpose()*(posDesFootWorld - eState->m_position)-quadKine->GetHipLocation(foot);
            Eigen::Vector3f velDesFootLeg = eState->m_rotation.transpose()*(velDesFootWorld - eState->m_velWorld);
            legsCommand[foot].m_footPosDesired = posDesFootLeg;
            legsCommand[foot].m_footVelDesired = velDesFootLeg;
            legsCommand[foot].m_kpCartesian = _kp;
            legsCommand[foot].m_kdCartesian = _kd;

        }
        // * Foot in stance state
        else{
//            std::cout<<"stance"<<std::endl;
            _firstSwing[foot] = true;
            Eigen::Vector3f posDesFootWorld = _swingPlanner[foot].GetPosition();
//            std::cout<<"Desired foot"<<foot<<" pos world:"<<posDesFootWorld.transpose()<<std::endl;
            Eigen::Vector3f velDesFootWorld = _swingPlanner[foot].GetVelocity();
            Eigen::Vector3f posDesFootLeg = eState->m_rotation.transpose()*(posDesFootWorld-eState->m_position)-quadKine->GetHipLocation(foot);
//            std::cout<<"Desired foot"<<foot<<" pos leg:"<<posDesFootLeg.transpose()<<std::endl;
            Eigen::Vector3f velDesFootLeg = eState->m_rotation.transpose()*(velDesFootWorld-eState->m_velWorld);
//            posDesFootLeg[2]=-0.25;
            legsCommand[foot].m_footPosDesired = posDesFootLeg;
            legsCommand[foot].m_footVelDesired = velDesFootLeg;
            legsCommand[foot].m_kdCartesian = _kdStance;
            legsCommand[foot].m_kpCartesian = _kpStance;
            legsCommand[foot].m_footForce = _footForce[foot];
//            legsCommand[foot].m_kdJoints = Eigen::Matrix3f::Identity()*0.2;
        }
    }
//    std::cout<<"Desired foot pos:"<<legsCommand[0].m_footPosDesired.transpose()<<std::endl;

}

void ConvexMPController::UpdateMPCIfNeeded(int *mpcTable, EstimatedState *const eState, bool omniMode) {
    if((_loopsCounter%_loopsPerMPC)==0){//(_loopsCounter%_loopsPerMPC)==0
        float* ePos=eState->m_position.data();
        Eigen::Vector3f velDesBody(_xVelDes,_yVelDes,0);
        Eigen::Vector3f velDesWorld=omniMode ? velDesBody : eState->m_rotation*velDesBody;
        // * Set desired state trajectory of the whole MPC horizon
        if(_currentGait==GaitName::standing){
            float trajInitial[12]={ _rollDes,_pitchDes,eState->m_rpy[2],                     //Posture
                                    0,0,_bodyHeight, //Position
                                    0,0,0,0,0,0};                                            //Velocity
            for(int i=0;i<_nSegment;i++){
                for(int j=0;j<12;j++){
                    _trajAll[12*i+j]=trajInitial[j];
                }
            }
        }
        else{
            const float maxPosError=.1;
            float xStart=_posDes[0];
            float yStart=_posDes[1];
            if(xStart-ePos[0]>maxPosError) xStart=ePos[0]+maxPosError;
            if(ePos[0]-xStart>maxPosError) xStart=ePos[0]-maxPosError;
            if(yStart-ePos[1]>maxPosError) yStart=ePos[1]+maxPosError;
            if(ePos[1]-yStart>maxPosError) yStart=ePos[1]-maxPosError;
            ePos[0]=xStart;
            ePos[1]=yStart;
            float trajInitial[12]={ 0.,0.,0.,
                                    0,0,_bodyHeight,
                                    0.,0.,0,
                                    0,0,0};
//            float trajInitial[12]={ 0.,0.,_yawDes,
//                                    xStart,yStart,_bodyHeight,
//                                    0.,0.,_yawVelDes,
//                                    velDesWorld[0],velDesWorld[1],0};
//            for(int i=0;i<12;i++)
//                std::cout<<trajInitial[i]<<" ";
//            std::cout<<std::endl;
//            std::cout<<_yawVelDes<<std::endl;
            for(int i=0;i<_nSegment;i++){
                for(int j=0;j<12;j++){
                    _trajAll[12*i+j]=trajInitial[j];
                }
                if(i==0){
                    _trajAll[2]=eState->m_rpy[2];
                }
                else{
                    _trajAll[12*i+3]=_trajAll[12*(i-1)+3]+_dtMPC*velDesWorld[0];
                    _trajAll[12*i+4]=_trajAll[12*(i-1)+4]+_dtMPC*velDesWorld[1];
                    _trajAll[12*i+2]=_trajAll[12*(i-1)+2]+_dtMPC*_yawVelDes;
                }
            }

        }
        SolveDenseMPC(mpcTable,eState);
    }
}

void ConvexMPController::SolveDenseMPC(int * mpcTable,EstimatedState*const eState) {
    float Q[12]={2.5, 2.5, 10, 100, 100, 100, 0, 0, 0.5, 0.2, 0.2, 0.1};;
    float yaw=eState->m_rpy[2];
    float* weights=Q;
    float alpha=1e-5;
    float* p=eState->m_position.data();
    float* v=eState->m_velWorld.data();
//    std::cout<<"Vel World:"<<eState->m_velWorld.transpose()<<std::endl;
    float* w=eState->m_omegaBody.data();
//    w[0]=0;w[1]=0;w[2]=0;
    float q[4];
    q[0]=eState->m_quat.w();
    q[1]=eState->m_quat.x();
    q[2]=eState->m_quat.y();
    q[3]=eState->m_quat.z();
    float r[12]={0.180753,0.191189,-0.187461,-0.198793,
    -0.11155,0.109807,-0.110056,0.110388,
    -0.26213,-0.263185,-0.263304,-0.261944 };
//    for(int i=0;i<12;i++){
//        r[i]=_posFootWorld[i%4][i/4]-eState->m_position[i/4];
//    }
//    std::cout<<std::endl;
//    std::cout<<eState->m_rotation<<std::endl;
//    std::cout<<eState->m_quat.coeffs().transpose()<<std::endl;
//    std::cout<<eState->m_rpy.transpose()<<std::endl;
    // * Set problem, update data, solve problem
    setup_problem(_dtMPC,_nSegment,0.4,120);
    update_solver_settings(10000,1e-07,1e-08,1.5,0.1,0);
    update_problem_data_floats(p,v,q,w,r,yaw,weights,_trajAll,alpha,mpcTable);
    // * Update output foot force
    for(int leg=0;leg<4;leg++){
        Eigen::Vector3f f;
        for(int axis = 0; axis < 3; axis++)
            f[axis] = get_solution(leg*3 + axis);
        _footForce[leg] = -eState->m_rotation.transpose()*f;
        std::cout << "Foot world" << leg << " force: " << f.transpose() << "\n";
//        std::cout << "Foot body" << leg << " force: " << _footForce[leg].transpose() << "\n";
    }
    std::cout<<std::endl;
}